import os
import random


def creatRandomDataTxt(train_local,data_path, training_rate, validation_rate):
    data_path = os.path.abspath(data_path)
    # get all the img paths in data
    # imgs_list=getImgFromOriginalDataSet(data_path)
    imgs_list = getImgFromProcessedDataSet(data_path)

    train_f_path = os.path.join('data','train.csv')
    test_f_path = os.path.join('data','test.csv')
    valid_f_path = os.path.join('data','valid.csv')
    if not train_local:
        train_f_path = os.path.join('LaneLineDetection', train_f_path)
        test_f_path = os.path.join('LaneLineDetection', test_f_path)
        valid_f_path = os.path.join('LaneLineDetection', valid_f_path)
    indexes = [i for i in range(len(imgs_list))]
    random.shuffle(indexes)
    num_training_valid_data = int(training_rate * len(imgs_list))
    num_valid_data = int(num_training_valid_data * validation_rate)
    num_training_data = num_training_valid_data - num_valid_data
    training_img_list = imgs_list[:num_training_data]
    validation_img_list = imgs_list[num_training_data:num_training_data + num_valid_data]
    testing_img_list = imgs_list[num_training_data + num_valid_data:]
    with open(train_f_path, 'w') as f:
        for i, img_path in enumerate(training_img_list):
            dirs = img_path.split(os.sep)
            # label_path = os.path.join(data_path,"Gray_Label", 'Label_'+dirs[-6].lower(),"Label",dirs[-3],dirs[-2],dirs[-1].replace(".jpg","_bin.png"))
            label_path = os.path.join(data_path, "label", dirs[-1].replace(".jpg", "_bin.png"))
            f.writelines(img_path + ',' + label_path + '\n')
    with open(valid_f_path, 'w') as f:
        for i, img_path in enumerate(validation_img_list):
            dirs = img_path.split(os.sep)
            # label_path = os.path.join(data_path, "Gray_Label", 'Label_' + dirs[-6].lower(), "Label", dirs[-3], dirs[-2],
            #                           dirs[-1].replace(".jpg", "_bin.png"))
            label_path = os.path.join(data_path, "label", dirs[-1].replace(".jpg", "_bin.png"))

            f.writelines(img_path + ',' + label_path + '\n')
    with open(test_f_path, 'w') as f:
        for i, img_path in enumerate(testing_img_list):
            dirs = img_path.split(os.sep)
            # label_path = os.path.join(data_path, "Gray_Label", 'Label_' + dirs[-6].lower(), "Label", dirs[-3], dirs[-2],
            #                           dirs[-1].replace(".jpg", "_bin.png"))
            label_path = os.path.join(data_path, "label", dirs[-1].replace(".jpg", "_bin.png"))
            f.writelines(img_path + ',' + label_path + '\n')


def getImgFromOriginalDataSet(data_path):
    imgs_list = []
    root_dirs = os.listdir(data_path)
    for dir in root_dirs:
        road_path = os.path.join(data_path, dir)
        if dir.startswith("Road") and os.path.isdir(road_path):
            color_img_path = os.path.join(road_path, "ColorImage_" + dir.lower(), "ColorImage")
            record_dirs = os.listdir(color_img_path)
            for record_dir in record_dirs:
                record_path = os.path.join(color_img_path, record_dir)
                camera_dirs = os.listdir(record_path)
                for camera_dir in camera_dirs:
                    camera_path = os.path.join(record_path, camera_dir)
                    imgs = os.listdir(camera_path)
                    for img in imgs:
                        img_path = os.path.join(camera_path, img)
                        if os.path.isfile(img_path):
                            imgs_list.append(img_path)


def getImgFromProcessedDataSet(data_path):
    imgs_list = []
    input_dir = os.listdir(os.path.join(data_path, "input"))
    for img_name in input_dir:
        img_path = os.path.join(data_path, "input", img_name)
        if os.path.isfile(img_path):
            imgs_list.append(img_path)
    return imgs_list


if __name__ == "__main__":
    # creatRandomDataTxt("../../RoadLineDataset", 0.8, 0.1)
    creatRandomDataTxt("../../RoadLineDataset/processed_data", 0.8, 0.1)
